feat(pipeline): run pipeline stages in parallel - #873
Open
albertlast wants to merge 1 commit into
Open
Conversation
Pages were processed strictly one stage at a time. This lets stages overlap and lets a stage fold several pages into one model call. Two config knobs control it, both `0` = auto: - `max_inflight_pages` — pages moving through the pipeline at once. Set to `1` to restore fully sequential behaviour. - `max_batch_pages` — pages a single stage folds into one model call. The larger VRAM lever, so lower this one first. Engines opt in rather than out. The `Engine` trait gains `max_workers`, `max_batch` and `run_batch`, all defaulting to no concurrency, so an engine that shares one GPU context or a `&mut` model is untouched. `run_batch` returns one result per input in order, so a failure is attributed to the page that caused it instead of failing the group. Supporting changes: - `Registry::get` dedupes concurrent misses behind a per-engine lock, so parallel stages hitting a cold engine no longer each load the model and allocate its GPU memory. - Stage threads are pooled and park between runs instead of exiting: candle caches cuDNN handles in a `thread_local!` that is unsafe to tear down. - Translation batching is skipped when a custom system prompt is set, since such a prompt describes the single-page `[N]` block format. - Settings UI for both knobs, translated into all nine locales. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
albertlast
requested review from
Map1en,
fffonion,
karrot0,
liksunrice and
mayocream
as code owners
July 25, 2026 17:34
Contributor
|
Thanks for your first PR to Koharu. Please review our contribution guide before review: In the PR description, include:
If AI helped produce the patch, a human still needs to review and understand it before submission. |
Author
|
when i see correctly a pr with similiar scoping got merged, so can i close this pr? @mayocream |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
I translate long manga — runs of 1000+ pages. At that scale the bottleneck isn't translation quality; the existing per-page quality is already good enough for what I need. It's wall-clock time. The pipeline currently runs one page through one stage at a time, so the GPU sits idle during OCR, the CPU sits idle during inference, and a long run takes far longer than the hardware requires.
This PR is aimed squarely at that: throughput on large batches, without changing what any single page produces.
What changed
Pages currently move through the pipeline strictly one stage at a time. This PR lets stages overlap, and lets a stage fold several pages into a single model call.
Two config knobs control it, both
0= auto:1meansmax_inflight_pagesmax_batch_pagesmax_batch_pagesis the larger VRAM lever, so it's the one to lower first.Which stage actually gets what
There are two independent wins here, and it's worth separating them:
1. Cross-stage overlap — applies to every stage, including the GPU-bound ones. Page 2 can run detection while page 1 is inpainting. No engine opts into this and none can opt out; it falls out of the streaming driver. For a 1000-page run this is the bulk of the gain.
2. Within-stage parallelism and batching — opt-in, per engine. The
Enginetrait gainsmax_workers,max_batchandrun_batch, all defaulting to no concurrency. Only 4 of the 15 registered engines override them, and deliberately asymmetrically:comic-text-detector(+-seg,comic-text-bubble-detector,anime-text)manga-ocrMangaOcr::inferenceis a genuine tensor batch — crops are cat'd into one forward passpaddle-ocr-vl-1.6,mit48px-ocryuzumarker-font-detectionFontDetector::inferencepreprocesses crops in parallel then does one batched forwardllm&mutand would just queue on its state locklama-manga,aot-inpainting,flux2-kleinkoharu-renderermin(cores, 4)&self; the only shared state is a short-lived font-book mutexrun_batchreturns one result per input, in order, so a failure is attributed to the page that caused it instead of failing the whole group.Supporting changes:
Registry::getnow dedupes concurrent misses behind a per-engine lock. Without it, parallel stages hitting a cold engine each load the model and allocate its GPU memory, then discard all but one.[N]block format and can't be assumed to teach the batched[bP-N]form.Stopping a run (and why stage threads are pooled)
Cancelling was where the naive version fell over, and it's what drove the least obvious part of this design.
The first cut gave each stage its own thread that exited when the stage ended. Pressing Stop ends every stage thread at once, and that turned out to crash the app in two separate ways:
candlecaches its cuDNN handles in athread_local!(neitherSendnorSync), andcudarc'sDrop for CudnnunwrapscudnnDestroy. A thread that ends therefore destroys those handles and turns any teardown error into a panic inside a destructor — a hard crash, and cancelling is the easiest way to hit it.Arc<ClientWithMiddleware>shared by every stage. Dropping a stage's runtime takes its pooled connections with it, so a different, still-running worker's next request fails witherror sending request.So stage threads now park instead of exiting, and are reused across runs. That is the whole reason the pool exists — it isn't premature optimisation.
Cancellation semantics themselves:
Itemitself, so dropping an item anywhere — success, failure, or cancellation — releases its slot with no explicit bookkeeping and no leak path.There are unit tests for the pool covering exactly this: that thread-locals survive across stages, that a panicking stage does not kill its pooled thread, and that a closed-and-drained channel ends the stage cleanly.
Known gap: MCP does not honour these settings
start_pipeline(the HTTP route the UI uses) reads the limits from app config, so the Settings knobs apply there. The MCP entry point still passesPipelineLimits::default()— i.e. auto — so a whole-project run driven over MCP will parallelise even if the user setmax_inflight_pages = 1. ThepipelineCLI binary does the same, though it's single-page so it makes no practical difference there.I left it as-is rather than guess at the intended layering: these limits are currently app config, not part of
StartPipelineRequest, so it isn't obvious whether MCP should read global config or whether the limits belong on the request instead (which is roughly what #830 does with its per-run flag). Happy to wire it up either way — just say which you'd prefer.Auto mode is not a black box
0means auto, and the Settings UI resolves and displays what auto actually picked on this machine rather than leaving the user guessing:The worker count depends on the host's core count, which the frontend can't know, so
MetaInfogains acpu_workersfield reported by the server (GET /meta). The effectivemax_inflight_pages/max_batch_pagesare exposed there too, so the displayed numbers are the ones the driver will really use — not a client-side guess.I want to flag this up front rather than have a maintainer discover it at merge time.
Page completion order becomes non-deterministic. That is inherent to this PR, not incidental: pages are streamed through stages and finish out of order, which is exactly why a
Tracker::frontierexists — progress is reported from the lowest unfinished page so the percentage stays monotonic even when page 3 finishes before page 1. There are tests asserting this (frontier_never_decreases_under_interleaved_completion).#830 wants the opposite for its chapter path. Its own settings string says it will "run detect and OCR on all pages first, then translate with shared chapter context", and it advertises that "translation preserves reading order across pages." That is a barrier plus a stable global ordering — the two properties this PR deliberately gives up on the default path.
Structurally the two collide in the same place. #830 splits
run()intorun_sequential()+run_chapter_mode()and inserts at@@ -149,6 +193,66 @@ pub async fn run(. This PR replaces the body of that samerun()with the streaming stage-worker driver. We also both touchpipeline/engine.rs,llm.rs,bin/pipeline.rs,koharu-llm/src/prompt.rs,rpc/mcp/mod.rs,rpc/routes/pipelines.rs,SettingsDialog.tsx,openapi.jsonand the locale files. Whichever lands second will need a real rebase, not a mechanical one.That said, I don't think they're fundamentally incompatible — and the two are arguably chasing different goals. #830 trades time for cross-page quality; this PR trades ordering for time. Its chapter mode is an early return into its own barriered code path, so it would keep its ordering guarantees regardless of what the default path does — the parallel driver simply wouldn't apply to it. The reconciliation is mostly mechanical-but-tedious rather than a design dead end, and
run_step_for_all_pagescould later use the same batching machinery.User-visible behavior
0is inspectable rather than opaque.max_inflight_pages = 1restores the previous sequential behavior exactly, ordering included.How I verified
cargo fmt -- --check,cargo check,cargo clippy -- -D warnings— clean.cargo test --workspace --tests— all pass.bun lint:uiandbun run --filter ui test(156 tests, 23 files) — all pass.bun run generate:openapireproduces the committedui/openapi.jsonwith no drift.0as a real "auto" value), and the thread pool — that thread-locals survive across stages, that a panicking stage doesn't kill its pooled thread, and that concurrent stages each get their own thread.I did not regenerate
tests/integration-tests/client. It's already out of sync withmain's ownopenapi.jsonand regenerating pulls in 250+ unrelated files from a different generator version, which would bury this diff. Happy to do it in a separate PR if you'd like.AI usage disclosure
Per the contributing guide: AI assistance (Claude Code) was used for parts of this patch. I have reviewed, run and understood the change, and I'm responsible for it. Happy to walk through any part of the design — particularly the thread-pool parking rationale and the ordering trade-off above, which are the least obvious pieces.
🤖 Generated with Claude Code